from flask import Flask, request, jsonify

from src.bert.bert_predict import bert_predict

app = Flask(__name__)


# @app.route('/bert', methods=['POST'])
def bert_api():
    data = request.get_json()
    # 从JSON数据中提取文本内容
    # 假设前端发送的数据是 {"text": "要分析的文本"}
    text = data.get('text', '') if isinstance(data, dict) else str(data)
    category, label = bert_predict(text)
    return jsonify({
        'category': category,
        'label': label
    })


if __name__ == '__main__':
    app.run(host='127.0.0.1', port=5000)
